import os
import struct
import threading
import hashlib
from typing import Dict, Any, Optional, Callable
from log_manager import LogManager
from error_handler import ErrorHandler
import time

class FileTransferManager:
    """文件传输管理器，负责处理文件的上传和下载。"""

    def __init__(self, protocol, device_manager):
        """初始化文件传输管理器。

        Args:
            protocol: 通信协议实例
            device_manager: 设备管理器实例
        """ 
        self.protocol = protocol
        self.device_manager = device_manager
        self.transfer_state = {
            'type': '',  # 'upload' 或 'download'
            'is_folder': False,  # 是否是文件夹传输
            'files': [],  # 文件列表（文件夹传输时使用）
            'current_index': 0,  # 当前传输的文件索引（文件夹传输时使用）
            'total_files': 0,  # 总文件数（文件夹传输时使用）
            'local_path': '',  # 本地路径
            'remote_path': '',  # 远程路径
            'size': 0,  # 当前文件大小
            'offset': 0,  # 当前文件已传输大小
            'start_time': 0,  # 开始时间
            'on_going': False,  # 是否正在传输
            'fp': None  # 文件句柄
        }
        self.write_buffer_size: int = 256 * 1024  # 256KB buffer
        self.read_buffer_size: int = 64 * 1024  # 64KB buffer
        self.data_buffer: bytearray = bytearray()
        self.transfer_callbacks: Dict[str, Callable] = {}
        self._is_transferring: bool = False
        self._lock: threading.Lock = threading.Lock()
        self.speed_limit: int = 0  # 0表示不限速，单位：字节/秒
        self.transfer_complete: threading.Event = threading.Event()
        self.transfer_result: bool = False

    def _calculate_md5(self, file_path, chunk_size=8192):
        """计算文件的MD5值。
        
        Args:
            file_path: 文件路径
            chunk_size: 每次读取的块大小
            
        Returns:
            str: MD5值的十六进制字符串，失败返回None
        """
        md5 = hashlib.md5()
        try:
            with open(file_path, 'rb') as f:
                while True:
                    chunk = f.read(chunk_size)
                    if not chunk:
                        break
                    md5.update(chunk)
            return md5.hexdigest()
        except Exception as e:
            LogManager.error(f"计算MD5失败: {str(e)}")
            return None

    def is_transferring(self):
        """检查是否正在传输文件。

        Returns:
            bool: 是否正在传输
        """
        return self._is_transferring

    def handle_response(self, cmd: int, data: bytes):
        """处理设备响应。

        Args:
            cmd: 命令
            data: 响应数据

        Returns:
            bool: 是否处理成功
        """
        try:
            if cmd == self.protocol.parser.CMD_WRITE_FILE_END:
                LogManager.debug(f"收到写文件结束响应，数据长度: {len(data)}")
                return self._handle_write_end_response(data)
            if cmd == self.protocol.parser.CMD_READ_FILE_DATA:
                return self._handle_read_data_response(data)
            LogManager.warning(f"未知的响应命令: {cmd}")
            return False
        except Exception as exc:
            LogManager.error(f"处理响应失败: {str(exc)}")
            return False

    def transfer_folder_to_device(self, local_dir: str, remote_dir: str):
        """文件夹上传方法"""
        try:
            # 创建文件夹传输线程
            folder_thread = threading.Thread(
                target=self._folder_upload_thread,
                args=(local_dir, remote_dir),
                name="FolderUploadThread"
            )
            folder_thread.daemon = True
            folder_thread.start()
            return True
            
        except Exception as exc:
            LogManager.error(f"启动文件夹上传失败: {str(exc)}")
            return False

    def _folder_upload_thread(self, local_dir: str, remote_dir: str):
        """文件夹上传线程"""
        try:
            # 获取本地文件夹结构
            files, folders = self._get_local_folder_structure(local_dir)
            
            # 在设备上创建文件夹结构
            self._create_remote_folder_structure(remote_dir, folders)
            
            # 初始化传输状态
            self.transfer_state.update({
                'type': 'upload',
                'is_folder': True,
                'files': files,
                'current_index': 0,
                'total_files': len(files),
                'local_path': local_dir,
                'remote_path': remote_dir,
                'on_going': True
            })
            
            # 逐个上传文件
            for index, file_info in enumerate(files):
                if not self.transfer_state.get('on_going', True):
                    raise Exception("传输已取消")
                    
                local_file = os.path.join(local_dir, file_info['path'])
                remote_file = os.path.join(remote_dir, file_info['path'])
                
                LogManager.info(f"开始上传第 {index + 1}/{len(files)} 个文件: {remote_file}")
                
                # 上传单个文件并等待完成
                self.transfer_file_to_device(local_file, remote_file, file_info['size'])
                self.transfer_complete.wait()
                
                if not self.transfer_result:
                    raise Exception(f"文件上传失败: {file_info['path']}")
                    
                self.transfer_state['current_index'] = index + 1
                
            LogManager.info("文件夹上传完成")
            
        except Exception as exc:
            LogManager.error(f"文件夹上传失败: {str(exc)}")
            self._notify_complete(False, str(exc))
            
        finally:
            self.transfer_state['on_going'] = False
            self._notify_complete(True)
            with self._lock:
                self._is_transferring = False

    def transfer_folder_from_device(self, remote_dir: str, local_dir: str):
        """文件夹下载方法"""
        try:
            # 创建文件夹传输线程
            folder_thread = threading.Thread(
                target=self._folder_download_thread,
                args=(remote_dir, local_dir),
                name="FolderDownloadThread"
            )
            folder_thread.daemon = True
            folder_thread.start()
            return True
            
        except Exception as exc:
            LogManager.error(f"启动文件夹下载失败: {str(exc)}")
            return False

    def _folder_download_thread(self, remote_dir: str, local_dir: str):
        """文件夹下载线程"""
        try:
            # 获取设备上的文件夹结构
            files, folders = self._get_device_folder_structure(remote_dir)
            
            # 在本地创建文件夹结构
            self._create_local_folder_structure(local_dir, folders)
            
            # 初始化传输状态
            self.transfer_state.update({
                'type': 'download',
                'is_folder': True,
                'files': files,
                'current_index': 0,
                'total_files': len(files),
                'local_path': local_dir,
                'remote_path': remote_dir,
                'on_going': True
            })
            
            # 逐个下载文件
            for index, file_info in enumerate(files):
                if not self.transfer_state.get('on_going', True):
                    raise Exception("传输已取消")
                    
                remote_file = os.path.join(remote_dir, file_info['path'])
                local_file = os.path.join(local_dir, file_info['path'])
                
                LogManager.info(f"开始下载第 {index + 1}/{len(files)} 个文件: {local_file}")
                
                # 下载单个文件并等待完成
                self.transfer_file_from_device(
                    remote_file, 
                    local_file, 
                    file_info['size']
                )
                self.transfer_complete.wait()
                
                if not self.transfer_result:
                    raise Exception(f"文件下载失败: {file_info['path']}")
                    
                self.transfer_state['current_index'] = index + 1
                
            LogManager.info("文件夹下载完成")
            
        except Exception as exc:
            LogManager.error(f"文件夹下载失败: {str(exc)}")
            self._notify_complete(False, str(exc))
            
        finally:
            self.transfer_state['on_going'] = False
            with self._lock:
                self._is_transferring = False
            self._notify_complete(True)

    def _get_local_folder_structure(self, local_dir: str) -> tuple[list, list]:
        """获取本地文件夹结构
        
        Args:
            local_dir: 本地文件夹路径
            
        Returns:
            tuple: (files, folders)
                - files: 文件列表，每个文件包含 path（相对路径）和 size 属性
                - folders: 文件夹列表，每个文件夹包含 path（相对路径）属性
        """
        files = []
        folders = []
        try:
            for root, dirs, filenames in os.walk(local_dir):
                # 获取相对路径
                rel_path = os.path.relpath(root, local_dir)
                if rel_path != '.':
                    # 添加文件夹
                    folders.append({
                        'path': rel_path
                    })
                
                # 添加文件
                for filename in filenames:
                    file_path = os.path.join(root, filename)
                    rel_file_path = os.path.relpath(file_path, local_dir)
                    files.append({
                        'path': rel_file_path,
                        'size': os.path.getsize(file_path)
                    })
                    
            return files, folders
        except Exception as e:
            LogManager.error(f"获取本地文件夹结构失败: {str(e)}")
            raise

    def _get_device_folder_structure(self, remote_dir: str) -> tuple[list, list]:
        """递归获取设备上的文件夹结构
        
        Args:
            remote_dir: 远程文件夹路径
            
        Returns:
            tuple: (files, folders)
                - files: 文件列表，每个文件包含 path 和 size 属性
                - folders: 文件夹列表，每个文件夹包含 path 属性
        """
        all_files = []
        all_folders = []
        
        def _recursive_list(current_dir: str, relative_base: str):
            """递归获取文件列表
            
            Args:
                current_dir: 当前要遍历的目录
                relative_base: 用于计算相对路径的基础目录
            """
            items = self.device_manager.list_device_files(current_dir)
            if items is None:
                raise Exception(f"获取设备文件列表失败: {current_dir}")
            
            for item in items:
                # 计算相对路径
                item_path = os.path.join(current_dir, item['name'])
                relative_path = os.path.relpath(item_path, relative_base)
                
                if item.get('is_directory'):
                    # 添加文件夹
                    all_folders.append({
                        'path': relative_path
                    })
                    # 递归处理子目录
                    _recursive_list(item_path, relative_base)
                else:
                    # 添加文件
                    all_files.append({
                        'path': relative_path,
                        'size': item['size']
                    })
        
        try:
            _recursive_list(remote_dir, remote_dir)
            return all_files, all_folders
        except Exception as e:
            LogManager.error(f"获取设备文件夹结构失败: {str(e)}")
            raise

    def _create_remote_folder_structure(self, remote_dir: str, folders: list):
        """在设备上创建文件夹结构
        
        Args:
            remote_dir: 远程文件夹路径
            folders: 文件夹列表，每个文件夹包含 path 属性
        """
        try:
            # 确保基础目录存在
            if not self.device_manager.create_directory(remote_dir):
                raise Exception(f"创建远程目录失败: {remote_dir}")
            
            # 在设备上创建文件夹
            for folder in folders:
                folder_path = os.path.join(remote_dir, folder['path'])
                if not self.device_manager.create_directory(folder_path):
                    raise Exception(f"创建远程目录失败: {folder_path}")
                
        except Exception as e:
            LogManager.error(f"创建远程文件夹结构失败: {str(e)}")
            raise

    def _create_local_folder_structure(self, local_dir: str, folders: list):
        """在本地创建文件夹结构
        
        Args:
            local_dir: 本地文件夹路径
            folders: 文件夹列表，每个文件夹包含 path 属性
        """
        try:
            # 确保基础目录存在
            os.makedirs(local_dir, exist_ok=True)
            
            # 创建文件夹
            for folder in folders:
                folder_path = os.path.join(local_dir, folder['path'])
                os.makedirs(folder_path, exist_ok=True)
                
        except Exception as e:
            LogManager.error(f"创建本地文件夹结构失败: {str(e)}")
            raise

    def transfer_file_to_device(self, local_path: str, remote_path: str, size: int):
        """传输文件到设备
        
        Args:
            local_path: 本地文件路径
            remote_path: 远程文件路径
            size: 文件大小
        """
        try:
            # 重置传输完成事件和结果
            self.transfer_complete.clear()
            self.transfer_result = False
            
            # 更新当前文件的状态
            self.transfer_state.update({
                'size': size,
                'offset': 0,
                'remote_path': remote_path,
                'local_path': local_path,
                'on_going': True
            })
            
            # 创建传输线程
            transfer_thread = threading.Thread(
                target=self._start_write_file_to_device,
                name="UploadThread"
            )
            transfer_thread.daemon = True
            transfer_thread.start()
            return True
            
        except Exception as e:
            ErrorHandler.handle_transfer_error(e)
            return False

    def transfer_file_from_device(self, remote_path: str, local_path: str, size: int):
        """从设备传输文件
        
        Args:
            remote_path: 远程文件路径
            local_path: 本地文件路径
            size: 文件大小
        """
        try:
            # 重置传输完成事件和结果
            self.transfer_complete.clear()
            self.transfer_result = False
            
            # 更新当前文件的状态
            self.transfer_state.update({
                'size': size,
                'offset': 0,
                'remote_path': remote_path,
                'local_path': local_path,
                'on_going': True
            })
            
            # 创建传输线程
            transfer_thread = threading.Thread(
                target=self._start_read_file_from_device,
                name="DownloadThread"
            )
            transfer_thread.daemon = True
            transfer_thread.start()
            return True
            
        except Exception as e:
            self._notify_complete(False, str(e))
            return False

    def _start_write_file_to_device(self):
        """开始上传文件"""
        try:
            # 发送开始命令
            path_bytes = self.transfer_state['remote_path'].encode('utf-8') + b'\0'
            LogManager.info(f"开始上传文件: {self.transfer_state['remote_path']}")
            
            self.protocol.send_command(
                self.protocol.parser.CMD_WRITE_FILE_START,
                path_bytes
            )
            
            # 等待响应
            if not self.protocol.wait_response(
                self.protocol.parser.CMD_WRITE_FILE_START):
                raise Exception("设备拒绝写入请求")

            self.transfer_state['on_going'] = True

            # 开始传输数据
            with open(self.transfer_state['local_path'], 'rb') as f:
                last_progress_time = time.time()
                
                while self.transfer_state['offset'] < self.transfer_state['size']:
                    chunk_size = min(
                        self.write_buffer_size,
                        self.transfer_state['size'] - self.transfer_state['offset']
                    )
                    data = f.read(chunk_size)
                    LogManager.debug(f"发送数据块: 偏移={self.transfer_state['offset']}, 大小={len(data)}")
                    # 发送数据块
                    content = struct.pack('<II', 
                        self.transfer_state['offset'],
                        len(data)
                    ) + data
                    
                    if not self.protocol.send_command(
                        self.protocol.parser.CMD_WRITE_FILE_DATA,
                        content
                    ):
                        raise Exception("发送数据失败")
                        
                    # 等待响应
                    if not self.protocol.wait_response(
                        self.protocol.parser.CMD_WRITE_FILE_DATA):
                        raise Exception("设备写入失败")
                        
                    self.transfer_state['offset'] += len(data)
                    
                    # 控制进度通知频率
                    now = time.time()
                    if now - last_progress_time >= 0.1:  # 每100ms更新一次
                        self._notify_progress(
                            self.transfer_state['offset'],
                            self.transfer_state['size']
                        )
                        last_progress_time = now

                    if self.transfer_state['on_going'] == False:
                        return False
                    
            # 发送结束命令
            self.protocol.send_command(self.protocol.parser.CMD_WRITE_FILE_END)

            # 等待MD5响应
            response = self.protocol.wait_response_data(
                self.protocol.parser.CMD_WRITE_FILE_END)
            if not response:
                raise Exception("获取文件MD5失败")
            
            # 处理MD5响应
            if len(response) >= 16:
                md5_hex = ''.join([f'{b:02x}' for b in response[:16]])
                LogManager.debug(f"解析到文件MD5: {md5_hex}")
                # 最后一次进度通知
                self._notify_progress(
                    self.transfer_state['size'],
                    self.transfer_state['size']
                )
                # 根据是否是文件夹内的文件显示不同的日志
                if 'is_folder_transfer' in self.transfer_state:
                    LogManager.info(f"文件夹内文件上传完成: {self.transfer_state['remote_path']}")
                else:
                    LogManager.info(f"文件上传完成: {self.transfer_state['remote_path']}")
                self._notify_complete(True, None, md5_hex)
                self.transfer_result = True
                return True
            else:
                raise Exception("无效的MD5响应数据")
            
        except Exception as e:
            LogManager.error(f"文件上传失败: {str(e)}")
            self._notify_complete(False, str(e))
            self.transfer_result = False
            return False
            
        finally:
            self.transfer_complete.set()  # 设置完成事件
            
    def _start_read_file_from_device(self):
        """开始从设备读取文件"""
        fp = None
        try:
            # 获取完整的远程和本地路径，确保远程路径使用正斜杠
            remote_path = self.transfer_state['remote_path'].replace('\\', '/')
            local_path = self.transfer_state['local_path']
                            
            LogManager.debug(f"准备下载文件:")
            LogManager.debug(f"  - 远程路径: {remote_path}")
            LogManager.debug(f"  - 本地路径: {local_path}")
                            
            path_bytes = remote_path.encode('utf-8') + b'\0'
            total_size = self.transfer_state['size']
            LogManager.info(f"开始从设备读文件: {remote_path}")
            LogManager.info(f"文件总大小: {total_size}")

            # 计算总块数和最后一块大小
            full_blocks = total_size // self.read_buffer_size
            last_block_size = total_size % self.read_buffer_size
            total_blocks = full_blocks + (1 if last_block_size > 0 else 0)
            
            LogManager.info(f"传输计划: {total_blocks}块, "
                           f"完整块={full_blocks}个({self.read_buffer_size}字节/块), "
                           f"最后块={last_block_size}字节")
            
            # 确保目标文件夹存在
            os.makedirs(os.path.dirname(local_path), exist_ok=True)
            
            # 创建本地文件
            LogManager.debug(f"准备创建本地文件: {local_path}")
            fp = open(local_path, 'wb')
            self.transfer_state['fp'] = fp
            last_progress_time = time.time()
            
            self.transfer_state['on_going'] = True
            
            # 读取数据
            for block in range(total_blocks):
                if not self.transfer_state.get('on_going', True):
                    raise Exception("传输已取消")
                    
                # 计算当前块的大小和偏移
                offset = block * self.read_buffer_size
                size = self.read_buffer_size if block < full_blocks else last_block_size
                
                LogManager.debug(f"读取第{block + 1}/{total_blocks}块: 偏移={offset}, 大小={size}")
                
                # 构造请求
                request = bytearray()
                request.extend(struct.pack('<I', offset))
                request.extend(struct.pack('<I', size))
                request.extend(path_bytes)
                
                LogManager.debug(f"构造请求: offset=0x{offset:08X}, size=0x{size:08X}, path={remote_path}")
                
                # 发送请求
                if not self.protocol.send_command(self.protocol.parser.CMD_READ_FILE, request):
                    raise Exception("发送读取请求失败")
                
                # 等待响应
                response = self.protocol.wait_response_data(self.protocol.parser.CMD_READ_FILE)
                if not response:
                    raise Exception("读取数据失败: 设备无响应")
                    
                # 处理响应数据
                if len(response) == 0:
                    raise Exception("接收到空数据块")
                if len(response) > size:
                    LogManager.warning(f"数据块大小超出预期: 期望<={size}, 实际={len(response)}")
                    response = response[:size]
                    
                LogManager.debug(f"收到响应数据: 长度={len(response)}")
                
                # 写入数据
                fp.write(response)
                self.transfer_state['offset'] = offset + len(response)
                
                # 控制进度通知频率
                now = time.time()
                if now - last_progress_time >= 0.1:
                    self._notify_progress(self.transfer_state['offset'], total_size)
                    last_progress_time = now
                    
            LogManager.debug(f"已下载: {self.transfer_state['offset']}/{total_size} 字节")
            
            # 验证总大小
            if self.transfer_state['offset'] != total_size:
                raise Exception(f"文件大小不匹配: 期望={total_size}, 实际={self.transfer_state['offset']}")
            
            LogManager.info(f"文件下载完成: {local_path}")
            self._notify_complete(True)
            self.transfer_result = True
            
        except Exception as e:
            LogManager.error(f"文件下载失败: {str(e)}")
            self._notify_complete(False, str(e))
            self.transfer_result = False
            
        finally:
            if fp:
                fp.close()
                self.transfer_state['fp'] = None
            self.transfer_complete.set()

    def register_callback(self, event_type, callback):
        """注册回调函数"""
        self.transfer_callbacks[event_type] = callback
        
    def _notify_progress(self, transferred: int, total: int):
        """统一的进度通知
        
        Args:
            transferred: 当前文件已传输大小
            total: 当前文件总大小
        """
        if 'progress' in self.transfer_callbacks:
            # 计算速度
            elapsed = time.time() - self.transfer_state.get('start_time', 0)
            speed = (transferred / elapsed) / 1024 if elapsed > 0 else 0  # KB/s
            
            # 生成状态信息
            if self.transfer_state.get('is_folder', False):
                current_file = self.transfer_state['current_index'] + 1
                total_files = self.transfer_state['total_files']
                
                # 计算总进度
                files = self.transfer_state['files']
                total_size = sum(f['size'] for f in files)
                transferred_size = sum(f['size'] for f in files[:self.transfer_state['current_index']])
                transferred_size += transferred  # 加上当前文件的已传输大小
                
                progress = (transferred_size / total_size * 100) if total_size > 0 else 0
                status = (f"文件 {current_file}/{total_files}, "
                         f"{speed:.1f}KB/s, {progress:.1f}%")
                
                # 传递总体进度
                self.transfer_callbacks['progress'](transferred_size, total_size, status)
            else:
                progress = (transferred / total * 100) if total > 0 else 0
                status = f"{speed:.1f}KB/s, {progress:.1f}%"
                self.transfer_callbacks['progress'](transferred, total, status)

    def _notify_complete(self, success: bool, error: Optional[str] = None, md5: Optional[str] = None):
        """通知传输完成
        
        Args:
            success: 是否成功
            error: 错误信息（如果失败）
            md5: 文件的MD5值（如果成功）
        """
        # 设置传输结果
        self.transfer_result = success
        
        # 如果是文件夹传输中的单个文件，只设置事件不触发回调
        if self.transfer_state.get('is_folder') and self.transfer_state.get('on_going'):
            return
        
        # 整个传输完成时才触发回调
        if 'complete' in self.transfer_callbacks and self.transfer_callbacks['complete']:
            self.transfer_callbacks['complete'](success, error, md5)
            self._is_transferring = False
        
  
  
    def cancel_transfer(self, silent=False):
        """取消传输
    
        Args:
            silent (bool): 是否静默取消，不触发回调
        """
        try:
            LogManager.info("正在取消传输...")
            
            if self.transfer_state.get('fp'):
                self.transfer_state['fp'].close()
                self.transfer_state['fp'] = None
                
            self.transfer_state['on_going'] = False
            
            
            # 确保设置传输状态为False
            with self._lock:
                self._is_transferring = False
            
            LogManager.info("传输已取消")
            
            # 只有在非静默模式下才通知取消
            if not silent:
                self._notify_complete(False, "传输已取消")
                
        except Exception as e:
            LogManager.error(f"取消传输失败: {str(e)}")
            if not silent:
                self._notify_complete(False, f"取消传输失败: {str(e)}")

    def transfer_to_device(self, local_path: str, remote_path: str, is_dir: bool = False, size: int = 0):
        """统一的上传接口
        
        Args:
            local_path: 本地路径
            remote_path: 远程路径
            is_dir: 是否是文件夹
            size: 文件大小
        """
        with self._lock:
            if self._is_transferring:
                raise Exception("已有传输任务正在进行")
            self._is_transferring = True
            
        try:
            # 检查本地路径是否存在
            if not os.path.exists(local_path):
                raise FileNotFoundError(f"本地路径不存在: {local_path}")
            
            # 重置传输状态
            self.transfer_state.update({
                'type': 'upload',
                'is_folder': is_dir,
                'local_path': local_path,
                'remote_path': remote_path,
                'on_going': True,
                'start_time': time.time(),
                'size': size
            })
            
            # 根据类型调用对应的传输方法
            if is_dir:
                return self.transfer_folder_to_device(local_path, remote_path)
            else:
                return self.transfer_file_to_device(local_path, remote_path, size)
            
        except Exception as e:
            LogManager.error(f"上传失败: {str(e)}")
            self._notify_complete(False, str(e))
            with self._lock:
                self._is_transferring = False
            return False

    def transfer_from_device(self, remote_path: str, local_path: str, is_dir: bool = False, size: int = 0):
        """统一的下载接口
        
        Args:
            remote_path: 远程路径
            local_path: 本地路径
            is_dir: 是否是文件夹
            size: 文件大小
        """
        with self._lock:
            if self._is_transferring:
                raise Exception("已有传输任务正在进行")
            self._is_transferring = True
            
        try:
            # 重置传输状态
            self.transfer_state.update({
                'type': 'download',
                'is_folder': is_dir,
                'local_path': local_path,
                'remote_path': remote_path,
                'on_going': True,
                'start_time': time.time(),
                'size': size
            })
            
            # 根据类型调用对应的传输方法
            if is_dir:
                return self.transfer_folder_from_device(remote_path, local_path)
            else:
                return self.transfer_file_from_device(remote_path, local_path, size)
            
        except Exception as e:
            LogManager.error(f"下载失败: {str(e)}")
            self._notify_complete(False, str(e))
            with self._lock:
                self._is_transferring = False
            return False
 